"""
MySQL MCP 服务器实现
提供基于 Model Context Protocol 的 MySQL 数据库访问能力
"""

import json
import logging
from typing import Dict, List, Optional, Union
import mysql.connector
from mysql.connector import Error

class MySQLMCPServer:
    """MySQL MCP 服务器
    
    实现 Model Context Protocol 规范的 MySQL 数据库访问服务器
    支持安全的数据库查询、模式检查和结果格式化
    """
    
    def __init__(self, config: Dict):
        """初始化 MySQL MCP 服务器
        
        Args:
            config: 数据库配置信息
                host: 数据库主机
                port: 端口
                user: 用户名
                password: 密码
                database: 数据库名
        """
        self.config = config
        self.connection = None
        self.logger = logging.getLogger(__name__)
        
    async def connect(self) -> None:
        """建立数据库连接"""
        try:
            self.connection = mysql.connector.connect(
                host=self.config["host"],
                port=self.config["port"],
                user=self.config["user"],
                password=self.config["password"],
                database=self.config["database"]
            )
            self.logger.info("MySQL connection established successfully")
        except Error as e:
            self.logger.error(f"Error connecting to MySQL: {e}")
            raise
            
    async def disconnect(self) -> None:
        """关闭数据库连接"""
        if self.connection and self.connection.is_connected():
            self.connection.close()
            self.logger.info("MySQL connection closed")
            
    async def handle_request(self, request: Dict) -> Dict:
        """处理 MCP 请求
        
        Args:
            request: MCP 请求
                action: 操作类型 (query, schema, metadata)
                query: SQL 查询语句
                parameters: 查询参数
                
        Returns:
            Dict: MCP 响应
                status: 状态码
                data: 响应数据
                metadata: 元数据
        """
        try:
            # 验证请求
            self._validate_request(request)
            
            # 根据操作类型处理
            action = request.get("action", "query")
            if action == "query":
                result = await self._handle_query(request)
            elif action == "schema":
                result = await self._get_schema()
            elif action == "metadata":
                result = await self._get_metadata()
            else:
                raise ValueError(f"Unsupported action: {action}")
                
            return {
                "status": "success",
                "data": result,
                "metadata": {
                    "server": "mysql",
                    "version": mysql.connector.__version__,
                    "database": self.config["database"]
                }
            }
            
        except Exception as e:
            self.logger.error(f"Error handling request: {e}")
            return {
                "status": "error",
                "error": str(e),
                "metadata": {
                    "server": "mysql"
                }
            }
            
    def _validate_request(self, request: Dict) -> None:
        """验证 MCP 请求
        
        Args:
            request: MCP 请求
            
        Raises:
            ValueError: 请求无效
        """
        if not isinstance(request, dict):
            raise ValueError("Request must be a dictionary")
            
        if "action" not in request:
            raise ValueError("Request must contain 'action' field")
            
        if request["action"] == "query":
            if "query" not in request:
                raise ValueError("Query request must contain 'query' field")
                
    async def _handle_query(self, request: Dict) -> List[Dict]:
        """处理查询请求
        
        Args:
            request: 查询请求
                query: SQL 查询语句
                parameters: 查询参数 (可选)
                
        Returns:
            List[Dict]: 查询结果
        """
        query = request["query"]
        parameters = request.get("parameters", {})
        
        # 验证和清理查询
        query = self._sanitize_query(query)
        
        try:
            cursor = self.connection.cursor(dictionary=True)
            cursor.execute(query, parameters)
            results = cursor.fetchall()
            cursor.close()
            
            return results
            
        except Error as e:
            self.logger.error(f"Error executing query: {e}")
            raise
            
    async def _get_schema(self) -> Dict:
        """获取数据库模式信息
        
        Returns:
            Dict: 数据库表和字段信息
        """
        schema = {}
        try:
            cursor = self.connection.cursor()
            
            # 获取所有表
            cursor.execute("""
                SELECT table_name 
                FROM information_schema.tables 
                WHERE table_schema = %s
            """, (self.config["database"],))
            
            tables = cursor.fetchall()
            
            # 获取每个表的字段信息
            for (table_name,) in tables:
                cursor.execute("""
                    SELECT column_name, data_type, is_nullable, column_key
                    FROM information_schema.columns
                    WHERE table_schema = %s AND table_name = %s
                """, (self.config["database"], table_name))
                
                columns = cursor.fetchall()
                schema[table_name] = [
                    {
                        "name": col[0],
                        "type": col[1],
                        "nullable": col[2] == "YES",
                        "key": col[3]
                    }
                    for col in columns
                ]
                
            cursor.close()
            return schema
            
        except Error as e:
            self.logger.error(f"Error getting schema: {e}")
            raise
            
    async def _get_metadata(self) -> Dict:
        """获取数据库元数据
        
        Returns:
            Dict: 数据库元数据信息
        """
        try:
            cursor = self.connection.cursor()
            
            # 获取数据库版本
            cursor.execute("SELECT VERSION()")
            version = cursor.fetchone()[0]
            
            # 获取数据库大小
            cursor.execute("""
                SELECT 
                    table_schema,
                    ROUND(SUM(data_length + index_length) / 1024 / 1024, 2) AS size_mb
                FROM information_schema.tables
                WHERE table_schema = %s
                GROUP BY table_schema
            """, (self.config["database"],))
            
            size = cursor.fetchone()[1]
            
            cursor.close()
            
            return {
                "version": version,
                "database": self.config["database"],
                "size_mb": size,
                "charset": self.connection.charset,
                "collation": self.connection.collation
            }
            
        except Error as e:
            self.logger.error(f"Error getting metadata: {e}")
            raise
            
    def _sanitize_query(self, query: str) -> str:
        """清理和验证 SQL 查询
        
        Args:
            query: SQL 查询语句
            
        Returns:
            str: 清理后的查询语句
            
        Raises:
            ValueError: 查询语句无效或不安全
        """
        # 移除多余的空白字符
        query = " ".join(query.split())
        
        # 检查是否是只读查询
        if not query.lower().startswith("select"):
            raise ValueError("Only SELECT queries are allowed")
            
        # TODO: 添加更多的安全检查
        
        return query 